﻿import paddle
import paddle.nn.functional as F
from paddle.vision.transforms import Compose, Normalize
import numpy as np
import os
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
import matplotlib.pyplot as plt
from tkinter import Checkbutton, IntVar
from PIL import Image, ImageDraw, ImageOps
import tkinter as tk

class AugmentedMNISTDataset(paddle.io.Dataset):
    def __init__(self, images, labels, transform=None):
        self.images = images
        self.labels = labels
        self.transform = transform

    def __getitem__(self, index):
        image = self.images[index]
        label = self.labels[index].astype('int64')  # 将标签转换为 int64
        if self.transform:
            image = self.transform(image)
        return image, label

    def __len__(self):
        return len(self.images)

# 定义数据预处理
transform = Compose([Normalize(mean=[127.5], std=[127.5], data_format='CHW')])

# 获取当前文件所在的目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 加载本地的扩展数据集
train_images = np.load(os.path.join(current_dir, 'augmented_train_images.npy'))
train_labels = np.load(os.path.join(current_dir, 'augmented_train_labels.npy'))
test_images = np.load(os.path.join(current_dir, 'augmented_test_images.npy'))
test_labels = np.load(os.path.join(current_dir, 'augmented_test_labels.npy'))

# 创建数据集对象
train_dataset = AugmentedMNISTDataset(train_images, train_labels, transform=transform)
test_dataset = AugmentedMNISTDataset(test_images, test_labels, transform=transform)

# 定义LeNet-5模型
class LeNetModel(paddle.nn.Layer):
    def __init__(self):
        super(LeNetModel, self).__init__()
        self.conv1 = paddle.nn.Conv2D(in_channels=1, out_channels=6, kernel_size=5, stride=1)  # 第一卷积层
        self.pool1 = paddle.nn.MaxPool2D(kernel_size=2, stride=2)  # 第一池化层
        self.conv2 = paddle.nn.Conv2D(in_channels=6, out_channels=16, kernel_size=5, stride=1)  # 第二卷积层
        self.pool2 = paddle.nn.MaxPool2D(kernel_size=2, stride=2)  # 第二池化层
        self.conv3 = paddle.nn.Conv2D(in_channels=16, out_channels=120, kernel_size=4, stride=1)  # 第三卷积层
        self.fc1 = paddle.nn.Linear(120, 84)  # 全连接层1
        self.fc2 = paddle.nn.Linear(84, 10)  # 输出层

    def forward(self, x):
        x = self.conv1(x)
        x = F.leaky_relu(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = F.leaky_relu(x)
        x = self.pool2(x)
        x = self.conv3(x)
        x = F.leaky_relu(x)
        x = paddle.flatten(x, start_axis=1, stop_axis=-1)  
        x = self.fc1(x)
        x = F.leaky_relu(x)
        out = self.fc2(x)
        return out

# 实例化模型
model = paddle.Model(LeNetModel())

# 准备模型
model.prepare(
    optimizer=paddle.optimizer.AdamW(parameters=model.parameters()),
    loss=paddle.nn.CrossEntropyLoss(),
    metrics=paddle.metric.Accuracy()
)

# 训练模型
model.fit(
    train_dataset,   #隐式地调用了getitem
    epochs=5,
    batch_size=16,
    verbose=1
)

# 评估模型
model.evaluate(test_dataset, verbose=1)

# 保存模型
model.save('output/mnist')

# 加载训练好的模型
model.load('output/mnist')

# 定义处理用户手写图片的函数
def process_image(image, invert=False):
    img = image.convert('L')  # 转换为灰度图像
    if invert:
        img = ImageOps.invert(img)  # 反转颜色
    img = img.resize((28, 28), Image.LANCZOS)  # 调整大小为28x28
    img = np.array(img).astype('float32')  # 转换为numpy数组
    img = img.reshape(1, 1, 28, 28)  # 调整形状为(1, 1, 28, 28)
    img = img / 127.5 - 1.0  # 标准化处理
    return img

def predict(image, model, invert=False):
    img = process_image(image, invert)
    result = model.predict_batch([img])
    predicted_label = np.argmax(result[0])
    return predicted_label, img.reshape(28, 28)

class DrawingBoard(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("Draw a digit")
        self.geometry("600x350")
        
        self.left_frame = tk.Frame(self, width=280, height=280, bg='white')
        self.left_frame.grid(row=0, column=0, padx=10, pady=10)

        self.canvas = tk.Canvas(self.left_frame, width=280, height=280, bg='white')
        self.canvas.pack()
        self.canvas.bind("<B1-Motion>", self.paint)
        
        self.right_frame = tk.Frame(self, width=280, height=280, bg='white')
        self.right_frame.grid(row=0, column=1, padx=10, pady=10)
        
        self.button_frame = tk.Frame(self)
        self.button_frame.grid(row=1, column=0, columnspan=2)

        self.predict_button = tk.Button(self.button_frame, text="Predict", command=self.predict)
        self.predict_button.pack(side=tk.LEFT)
        
        self.clear_button = tk.Button(self.button_frame, text="Clear", command=self.clear)
        self.clear_button.pack(side=tk.LEFT)

        self.result_label = tk.Label(self.button_frame, text="Predicted label: None")
        self.result_label.pack(side=tk.LEFT)

        self.invert_var = IntVar()
        self.invert_check = Checkbutton(self.button_frame, text="Invert Colors", variable=self.invert_var)
        self.invert_check.pack(side=tk.LEFT)

        self.fig, self.ax = plt.subplots(figsize=(2.8, 2.8))
        self.canvas_fig = FigureCanvasTkAgg(self.fig, master=self.right_frame)
        self.canvas_fig.get_tk_widget().pack()

        self.bg_color = 'white'
        self.fg_color = 'black'
        self.image = Image.new("L", (280, 280), 255)
        self.draw = ImageDraw.Draw(self.image)

    def paint(self, event):
        x1, y1 = (event.x - 10), (event.y - 10)
        x2, y2 = (event.x + 10), (event.y + 10)
        self.canvas.create_oval(x1, y1, x2, y2, fill=self.fg_color, width=20)
        self.draw.ellipse([x1, y1, x2, y2], fill=self.fg_color)

    def predict(self):
        invert = self.invert_var.get() == 1
        if self.bg_color == 'black':
            inverted_image = ImageOps.invert(self.image)
            label, img_28x28 = predict(inverted_image, model, invert=invert)
        else:
            label, img_28x28 = predict(self.image, model, invert=invert)
        
        self.result_label.config(text=f"Predicted label: {label}")
        
        # 在右侧显示28x28图片
        self.ax.clear()
        self.ax.imshow(img_28x28, cmap='gray')
        self.canvas_fig.draw()

    def clear(self):
        self.canvas.delete("all")
        self.image = Image.new("L", (280, 280), 255 if self.bg_color == 'white' else 0)
        self.draw = ImageDraw.Draw(self.image)
        self.result_label.config(text="Predicted label: None")
        self.ax.clear()
        self.canvas_fig.draw()

if __name__ == "__main__":
    app = DrawingBoard()
    app.mainloop()

